package pl.us.ij.plugins.segmentation.commands.watershed;

import ij.ImagePlus;
import ij.process.ColorProcessor;
import ij.process.FloatProcessor;
import ij.process.ImageProcessor;
import pl.us.ij.plugins.segmentation.Trace;
import pl.us.ij.plugins.segmentation.commands.ArrayHelper;
import pl.us.ij.plugins.segmentation.infrastructure.Command;
import pl.us.ij.plugins.segmentation.infrastructure.CommandParams;
import pl.us.ij.plugins.segmentation.infrastructure.ConvertStrategy;
import pl.us.ij.plugins.segmentation.infrastructure.GrayedColorConvertStrategy;
import pl.us.ij.plugins.segmentation.infrastructure.exceptions.BreakCommandException;

/**
 * Created by IntelliJ IDEA.
 * User: Przemek
 * Date: 2010-01-10
 * Time: 21:10:44
 * To change this template use File | Settings | File Templates.
 */
public class WatershedSegmentationCommand extends Command<WatershedSegmentationCommandParams> {
    private class Pixel {
        public float gray;
        public float gray2;
        public int x;
        public int y;
    }

    ImagePlus outImage;

    @Override
    protected ConvertStrategy prepareConvertStrategy(ImagePlus imp, CommandParams params) {
        ImageProcessor ip = imp.getProcessor();

        if ( ip instanceof ColorProcessor) {
            return new GrayedColorConvertStrategy(imp, params);
        }

        return super.prepareConvertStrategy(imp, params);
    }

    @Override
    protected ImagePlus prepareOutImage(ImagePlus inImage) {

        this.outImage = inImage.createImagePlus();
        this.outImage.setProcessor( new ColorProcessor(inImage.getWidth(), inImage.getHeight()) );

        return super.prepareOutImage(inImage);
    }

    @Override
    protected boolean finalizeOutImage(ImagePlus outImp) {
        this.outImage.updateAndDraw();
        this.outImage.show(createOutImageTitle());

        return false;
    }

    @Override
    protected void processImage(FloatProcessor fp) throws BreakCommandException {

        int width = fp.getWidth();
        int height = fp.getHeight();

        ImageProcessor ip = fp;

        float[] pixels = (float[]) ip.getPixels();
        Pixel[] pixs = new Pixel[pixels.length];
        Pixel[] pixs2 = new Pixel[pixels.length];

        int i = 0;
        for (int y = 0; y < height; y++) {
            for (int x = 0; x < width; x++) {
                pixs[i] = new Pixel();
                pixs[i].gray = pixels[i];
                pixs[i].x = x;
                pixs[i].y = y;
                i++;
            }

        }

        quick(0, pixs.length - 1, pixs);

        i = 0;
        for (int y = 0; y < height; y++) {
            for (int x = 0; x < width; x++) {
                pixs2[i] = new Pixel();
                pixs2[i].gray = -1;
                pixs2[i].gray2 = pixels[i];
                pixs2[i].x = x;
                pixs2[i].y = y;
                i++;
            }

        }

        int reGroupCount = 0;
        float k1 = -1;
        float k2 = -1;
        float color = 1;

        for (i = 0; i < pixels.length - 1; i++) {
            k1 = pixs[i].gray;
            float[] neighbors = new float[8];

            if (reGroupCount <= 0)
                reGroupCount = regroup(i, pixs, pixs2, width, height);
            reGroupCount--;

            int x = pixs[i].x;
            int y = pixs[i].y;

            if (pixs2[x + y * width].gray == -1)
                if (!((x == width - 1) || (y == height - 1) || (x == 0) || (y == 0))) {
                    neighbors[0] = pixs2[(x + y * width) - width - 1].gray;
                    neighbors[1] = pixs2[(x + y * width) - width].gray;
                    neighbors[2] = pixs2[(x + y * width) - width + 1].gray;
                    neighbors[3] = pixs2[(x + y * width) - 1].gray;
                    neighbors[4] = pixs2[(x + y * width) + 1].gray;
                    neighbors[5] = pixs2[(x + y * width) + width - 1].gray;
                    neighbors[6] = pixs2[(x + y * width) + width].gray;
                    neighbors[7] = pixs2[(x + y * width) + width + 1].gray;

                    boolean one = false;
                    boolean two = false;

                    float pool = 0;

                    for (int k = 0; k < 8; k++) {
                        if ((neighbors[k] != -1) && (neighbors[k] != 255))
                            if (!one) {
                                pool = neighbors[k];
                                one = true;
                            } else if (pool != neighbors[k]) {
                                two = true;
                            }

                    }

                    int shiftD = y * width;

                    if (!one && !two) {
                        if (color < 185)
                            color += 11;
                        else
                            color -= 185;

                        pixs2[x + y * width].gray = color;
                        pixels[shiftD + x] = color;
                    } else if (one && !two) {
                        pixs2[x + y * width].gray = pool;
                        pixels[shiftD + x] = pool;
                    } else if (one && two) {
                        pixs2[x + y * width].gray = 255;
                        pixels[shiftD + x] = 255;
                    }
                }

            if (k1 != k2) {
                k2 = k1;
            }
        }

        int shiftD = 0;
        int shift2 = width * (height - 1);

        for (int x = 0; x < width; x++) {
            pixels[shiftD + x] = 255;
            pixels[shift2 + x] = 255;
        }

        for (int y = 0; y < height; y++) {
            pixels[y * width] = 255;
            pixels[y * width + width - 1] = 255;
        }

        coloring2(width, height, pixels, 255);

        ip.setPixels(pixels);
    }

    public static void coloring(int width, int height, float[] pixels, float excludedColor) {

        final int maxPixel = (int) ArrayHelper.max(pixels);
        final int minPixel = (int) ArrayHelper.min(pixels);
        final int RANGE = maxPixel - minPixel + 1;
        final int SHIFT = -minPixel;

        float[][] rgbMap = new float[RANGE][3];
        for (int i = 0; i < RANGE; i++) {
            rgbMap[i][0] = -1;
            rgbMap[i][1] = -1;
            rgbMap[i][2] = -1;
        }
        int shiftD = 0;
        for (int y = 0; y < height; y++) {
            shiftD = y * width;

            for (int x = 0; x < width; x++) {
                float pixel = pixels[shiftD + x];

                if (pixel != excludedColor) {
                    if (rgbMap[(int) pixel + SHIFT][0] == -1) {
                        rgbMap[(int) pixel + SHIFT][0] = (float) Math.round(Math.random() * 200);
                        rgbMap[(int) pixel + SHIFT][1] = (float) Math.round(Math.random() * 200);
                        rgbMap[(int) pixel + SHIFT][2] = (float) Math.round(Math.random() * 200);
                    }
                }

                pixels[shiftD + x] = rgbMap[((int) pixel) + SHIFT][0];

            }
        }
    }

    public void coloring2(int width, int height, float[] pixels, float excludedColor) {

        final int maxPixel = (int) ArrayHelper.max(pixels);
        final int minPixel = (int) ArrayHelper.min(pixels);
        final int RANGE = maxPixel - minPixel + 1;
        final int SHIFT = -minPixel;

        ColorProcessor cp = (ColorProcessor) this.outImage.getProcessor();
        int[] outPixels = (int[])cp.getPixels();

        int[][] rgbMap = new int[RANGE][3];

        for (int i = 0; i < RANGE; i++) {
            rgbMap[i][0] = -1;
            rgbMap[i][1] = -1;
            rgbMap[i][2] = -1;
        }
        int shiftD = 0;
        for (int y = 0; y < height; y++) {
            shiftD = y * width;

            for (int x = 0; x < width; x++) {
                float pixel = pixels[shiftD + x];

                if (pixel != excludedColor) {
                    if (rgbMap[(int) pixel + SHIFT][0] == -1) {
                        rgbMap[(int) pixel + SHIFT][0] = (int) Math.round(Math.random() * 200);
                        rgbMap[(int) pixel + SHIFT][1] = (int) Math.round(Math.random() * 200);
                        rgbMap[(int) pixel + SHIFT][2] = (int) Math.round(Math.random() * 200);
                    }
                }

                outPixels[ shiftD + x ] = (rgbMap[(int) pixel + SHIFT][0] |
                                           rgbMap[(int) pixel + SHIFT][1] << 8 |
                                           rgbMap[(int) pixel + SHIFT][2] << 16);

            }
        }
    }

    public static void coloring(int width, int height, float[] pixels, float color, int channelNumber) {

        float[][] rgbMap = new float[256][3];
        for (int i = 0; i < 256; i++) {
            rgbMap[i][0] = -1;
            rgbMap[i][1] = -1;
            rgbMap[i][2] = -1;
        }
        int shiftD = 0;
        for (int y = 0; y < height; y++) {
            shiftD = y * width;

            for (int x = 0; x < width; x++) {
                float pixel = pixels[shiftD + x];

                if (pixel != color) {
                    if (rgbMap[(int) pixel][0] == -1) {
                        rgbMap[(int) pixel][0] = (float) Math.round(Math.random() * 200);
                        rgbMap[(int) pixel][1] = (float) Math.round(Math.random() * 200);
                        rgbMap[(int) pixel][2] = (float) Math.round(Math.random() * 200);
                    }
                }

                pixels[shiftD + x] = rgbMap[((int) pixel)][channelNumber];

            }
        }
    }

    public static void coloring(int width, int height, float[] pixels, int color, ImageProcessor ip) {

        final int maxPixel = (int) ArrayHelper.max(pixels);
        final int minPixel = (int) ArrayHelper.min(pixels);
        final int RANGE = maxPixel - minPixel + 1;
        final int SHIFT = -minPixel;

        float[][] rgbMap = new float[RANGE][3];

        FloatProcessor fim = ip.toFloat(0, null);
        float[] redPixels = (float[]) fim.getPixels();
        float[] greenPixels = (float[]) ip.toFloat(1, null).getPixels();
        float[] bluePixels = (float[]) ip.toFloat(2, null).getPixels();

        for (int i = 0; i < RANGE; i++) {
            rgbMap[i][0] = -1;
            rgbMap[i][1] = -1;
            rgbMap[i][2] = -1;
        }

        rgbMap[RANGE - 1][0] = 255;
        rgbMap[RANGE - 1][1] = 255;
        rgbMap[RANGE - 1][2] = 255;

        int shiftD = 0;
        for (int y = 0; y < height; y++) {
            shiftD = y * width;

            for (int x = 0; x < width; x++) {
                float pixel = pixels[shiftD + x];
                try {
                    if (pixel != color) {
                        if (rgbMap[(int) pixel + SHIFT][0] == -1) {
                            rgbMap[(int) pixel + SHIFT][0] = (float) Math.round(Math.random() * 200);
                            rgbMap[(int) pixel + SHIFT][1] = (float) Math.round(Math.random() * 200);
                            rgbMap[(int) pixel + SHIFT][2] = (float) Math.round(Math.random() * 200);
                        }
                    }
                }
                catch (Exception e) {
                    Trace.log("SHIFT: " + SHIFT);
                    Trace.log("RANGE: " + RANGE);
                    Trace.log("Pixel: " + pixel);
                    Trace.log("MinPixel: " + minPixel);
                    Trace.log("MaxPixel: " + maxPixel);
                }

                redPixels[shiftD + x] = rgbMap[((int) pixel) + SHIFT][0];
                greenPixels[shiftD + x] = rgbMap[((int) pixel) + SHIFT][1];
                bluePixels[shiftD + x] = rgbMap[((int) pixel) + SHIFT][2];
            }
        }
        ip.setPixels(0, fim);
        fim.setPixels(greenPixels);
        ip.setPixels(1, fim);
        fim.setPixels(bluePixels);
        ip.setPixels(2, fim);
    }

    private int regroup(int i, Pixel[] pixs, Pixel[] pixs2, int width, int height) {
        float bright = pixs[i].gray;
        int start = i;
        int finish = i;

        float[] neighbors = new float[8];
        int convert = 0;

        while (finish < pixs.length && pixs[finish].gray == bright) {
            int x = pixs[finish].x;
            int y = pixs[finish].y;

            if (!((x == width - 1) || (y == height - 1) || (x == 0) || (y == 0))) {
                neighbors[0] = pixs2[(x + y * width) - width - 1].gray;
                neighbors[1] = pixs2[(x + y * width) - width].gray;
                neighbors[2] = pixs2[(x + y * width) - width + 1].gray;
                neighbors[3] = pixs2[(x + y * width) - 1].gray;
                neighbors[4] = pixs2[(x + y * width) + 1].gray;
                neighbors[5] = pixs2[(x + y * width) + width - 1].gray;
                neighbors[6] = pixs2[(x + y * width) + width].gray;
                neighbors[7] = pixs2[(x + y * width) + width + 1].gray;

                for (int k = 0; k < 8; k++) {
                    if (neighbors[k] >= 0 && neighbors[k] < 255) {
                        convert = 1;
                        break;
                    }
                }

                if (convert == 1) {
                    Pixel pix = pixs[start];
                    pixs[start] = pixs[finish];
                    pixs[finish] = pix;
                    start++;
                    convert = 0;
                }
            }

            finish++;
        }

        return start - i == 0
                ? 1
                : start - i;
    }

    private void quick(int l, int p, Pixel[] pixs) {
        float v = pixs[(l + p) / 2].gray;
        int i = l;
        int j = p;

        do {
            while (pixs[i].gray < v)
                i++;
            while (v < pixs[j].gray)
                j--;

            if (i <= j) {
                Pixel pix = pixs[i];
                pixs[i] = pixs[j];
                pixs[j] = pix;
                i++;
                j--;
            }
        }
        while (i <= j);

        if (l < j)
            quick(l, j, pixs);
        if (i < p)
            quick(i, p, pixs);

//        Arrays.sort(pixs, new Comparator<Pixel>() {
//            public int compare(Pixel o1, Pixel o2) {
//                return o1.gray > o2.gray ? 1 : (o1.gray < o2.gray ? -1 : 0);
//            }
//        });
    }
}
